#include <png.h>
#include <string.h>
#include <iostream>
#include <memory>
#include <opencv2/opencv.hpp>
#include <vector>
#include "rga/RgaUtils.h"
#include "rga/im2d.hpp"
#include "rga/rga.h"
#include "time-statist.h"

using namespace std;

#define LOGINFO(fmt, args...) fprintf(stdout, "[INFO]%d: " fmt "\n", __LINE__, ##args)
#define LOGWARN(fmt, args...) fprintf(stdout, "[WARN]%d: " fmt "\n", __LINE__, ##args)
#define LOGERROR(fmt, args...) fprintf(stderr, "[ERROR]%d: " fmt "\n", __LINE__, ##args)

// 读取PNG文件并转换为RGB数据
std::vector<unsigned char> readPNGtoRGB(const std::string& filename, int& width, int& height) {
    std::vector<unsigned char> image_data;
    FILE* fp = fopen(filename.c_str(), "rb");
    if (!fp) {
        std::cerr << "Can't open file " << filename << std::endl;
        return image_data;
    }

    png_structp png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
    if (!png_ptr) {
        fclose(fp);
        return image_data;
    }

    png_infop info_ptr = png_create_info_struct(png_ptr);
    if (!info_ptr) {
        png_destroy_read_struct(&png_ptr, nullptr, nullptr);
        fclose(fp);
        return image_data;
    }

    if (setjmp(png_jmpbuf(png_ptr))) {
        png_destroy_read_struct(&png_ptr, &info_ptr, nullptr);
        fclose(fp);
        return image_data;
    }

    png_init_io(png_ptr, fp);
    png_read_info(png_ptr, info_ptr);

    width = png_get_image_width(png_ptr, info_ptr);
    height = png_get_image_height(png_ptr, info_ptr);
    int bit_depth = png_get_bit_depth(png_ptr, info_ptr);
    int color_type = png_get_color_type(png_ptr, info_ptr);

    // 确保是RGB图像
    if (color_type != PNG_COLOR_TYPE_RGB && color_type != PNG_COLOR_TYPE_RGBA) {
        png_destroy_read_struct(&png_ptr, &info_ptr, nullptr);
        fclose(fp);
        std::cerr << "Image is not in RGB or RGBA format." << std::endl;
        return image_data;
    }

    // 如果是RGBA，需要剥离alpha通道
    if (color_type == PNG_COLOR_TYPE_RGBA) {
        png_set_strip_alpha(png_ptr);
    }

    // 设置行滤波器
    png_read_update_info(png_ptr, info_ptr);

    image_data.resize(width * height * 3);  // RGB数据大小

    // 读取图像数据
    for (int y = 0; y < height; ++y) {
        png_bytep row = image_data.data() + y * width * 3;
        png_read_row(png_ptr, row, nullptr);
    }

    png_read_end(png_ptr, nullptr);
    png_destroy_read_struct(&png_ptr, &info_ptr, nullptr);
    fclose(fp);

    return image_data;
}

class CvtNV21ToRGB {
   public:
    CvtNV21ToRGB(int height, int width)
        : height_(height), width_(width) {
        srcSize_ = width * height * 3 / 2;
        dstSize_ = width * height * 3;
        src_ = new unsigned char[srcSize_];
        dst_ = new unsigned char[dstSize_];
        memset(&srcImg_, 0, sizeof(rga_buffer_t));
        memset(&dstImg_, 0, sizeof(rga_buffer_t));
        srcHandle_ = importbuffer_virtualaddr(src_, srcSize_);
        dstHandle_ = importbuffer_virtualaddr(dst_, dstSize_);
        if (srcHandle_ == 0 || dstHandle_ == 0) {
            printf("importbuffer failed!\n");
            throw std::bad_cast();
        }
        srcImg_ = wrapbuffer_handle(srcHandle_, width, height, srcFormat_);
        dstImg_ = wrapbuffer_handle(dstHandle_, width, height, dstFormat_);
        int ret = imcheck(srcImg_, dstImg_, {}, {});
        if (IM_STATUS_NOERROR != ret) {
            printf("%d, check error! %s", __LINE__, imStrError((IM_STATUS)ret));
            throw std::bad_cast();
        }
    }
    ~CvtNV21ToRGB() {
        if (srcHandle_) {
            releasebuffer_handle(srcHandle_);
        }
        if (dstHandle_) {
            releasebuffer_handle(dstHandle_);
        }
        if (src_) {
            delete src_;
            src_ = nullptr;
        }
        if (dst_) {
            delete dst_;
            dst_ = nullptr;
        }
    }
    void cvt2RGB(unsigned char* srcbuf) {
        memset(src_, 0, srcSize_);
        memset(dst_, 0, dstSize_);
        memcpy(src_, srcbuf, srcSize_);
        int ret = imcvtcolor(srcImg_, dstImg_, srcFormat_, dstFormat_);

        if (ret == IM_STATUS_SUCCESS) {
            // memcpy(dstbuf, dst_, dstSize_);
            // LOGINFO("cvtNV21ToRGB success");
        } else {
            LOGINFO("cvtNV21ToRGB running failed");
        }
    }

   public:
    int height_;
    int width_;
    int srcSize_, dstSize_;
    unsigned char *src_ = nullptr, *dst_ = nullptr;

   private:
    rga_buffer_t srcImg_, dstImg_;
    rga_buffer_handle_t srcHandle_, dstHandle_;

    int srcFormat_ = RK_FORMAT_YCbCr_420_SP;
    int dstFormat_ = RK_FORMAT_RGB_888;
};
void cvtNV21ToRGB(unsigned char* srcbuf, uint8_t*& dstbuf, int width, int height) {
    static bool init = true;
    static CvtNV21ToRGB* transor = new CvtNV21ToRGB(height, width);
    transor->cvt2RGB(srcbuf);
    dstbuf = transor->dst_;
}
// void cvtNV21ToRGB(unsigned char* srcbuf, unsigned char* dstbuf, int width, int height) {
//     int srcbufSize, dstbufSize, ret;
//     rga_buffer_t srcImg, dstImg;
//     rga_buffer_handle_t srcHandle, dstHandle;
//     bool success = false;

//     width = (width >> 4) << 4;
//     height = (height >> 2) << 2;
//     memset(&srcImg, 0, sizeof(srcImg));
//     memset(&dstImg, 0, sizeof(dstImg));
//     int srcFormat = RK_FORMAT_YCbCr_420_SP;
//     int dstFormat = RK_FORMAT_RGB_888;
//     srcbufSize = width * height * get_bpp_from_format(srcFormat);
//     dstbufSize = width * height * get_bpp_from_format(dstFormat);
//     memset(dstbuf, 0x80, dstbufSize);

//     srcHandle = importbuffer_virtualaddr(srcbuf, srcbufSize);
//     dstHandle = importbuffer_virtualaddr(dstbuf, dstbufSize);
//     if (srcHandle == 0 || dstHandle == 0) {
//         printf("importbuffer failed!\n");
//         goto release_buffer;
//     }

//     srcImg = wrapbuffer_handle(srcHandle, width, height, srcFormat);
//     dstImg = wrapbuffer_handle(dstHandle, width, height, dstFormat);

//     ret = imcheck(srcImg, dstImg, {}, {});
//     if (IM_STATUS_NOERROR != ret) {
//         printf("%d, check error! %s", __LINE__, imStrError((IM_STATUS)ret));
//         goto release_buffer;
//     }
//     TIMESTATIS_START(rga222rgb);
//     ret = imcvtcolor(srcImg, dstImg, srcFormat, dstFormat);
//     TIMESTATIS_END(rga222rgb);
//     if (ret == IM_STATUS_SUCCESS) {
//         success = true;
//     } else {
//         LOGINFO("cvtNV21ToRGB running failed");
//     }
// release_buffer:
//     if (srcHandle) {
//         releasebuffer_handle(srcHandle);
//     }
//     if (dstHandle) {
//         releasebuffer_handle(dstHandle);
//     }
//     if (!success) {
//         throw std::bad_cast();
//     }
// }

void cvtRGBToNV21(unsigned char* srcbuf, unsigned char* dstbuf, int width, int height) {
    int srcbufSize, dstbufSize, ret;
    rga_buffer_t srcImg, dstImg;
    rga_buffer_handle_t srcHandle, dstHandle;
    bool success = false;

    width = (width >> 4) << 4;
    height = (height >> 2) << 2;
    memset(&srcImg, 0, sizeof(srcImg));
    memset(&dstImg, 0, sizeof(dstImg));
    int srcFormat = RK_FORMAT_RGB_888;
    int dstFormat = RK_FORMAT_YCbCr_420_SP;
    srcbufSize = width * height * get_bpp_from_format(srcFormat);
    dstbufSize = width * height * get_bpp_from_format(dstFormat);
    memset(dstbuf, 0x80, dstbufSize);

    srcHandle = importbuffer_virtualaddr(srcbuf, srcbufSize);
    dstHandle = importbuffer_virtualaddr(dstbuf, dstbufSize);
    if (srcHandle == 0 || dstHandle == 0) {
        printf("importbuffer failed!\n");
        goto release_buffer;
    }

    srcImg = wrapbuffer_handle(srcHandle, width, height, srcFormat);
    dstImg = wrapbuffer_handle(dstHandle, width, height, dstFormat);

    ret = imcheck(srcImg, dstImg, {}, {});
    if (IM_STATUS_NOERROR != ret) {
        printf("%d, check error! %s\n", __LINE__, imStrError((IM_STATUS)ret));
        goto release_buffer;
    }
    TIMESTATIS_START(rga22yuv);
    ret = imcvtcolor(srcImg, dstImg, srcFormat, dstFormat);
    TIMESTATIS_END(rga22yuv);
    if (ret == IM_STATUS_SUCCESS) {
        success = true;
    } else {
        LOGINFO("cvtNV21ToRGB running failed");
    }
release_buffer:
    if (srcHandle) {
        releasebuffer_handle(srcHandle);
    }
    if (dstHandle) {
        releasebuffer_handle(dstHandle);
    }
    if (!success) {
        throw std::bad_cast();
    }
}
// 图片宽16对齐，高2对齐

int main(int argc, char* argv[]) {
    std::string filename = argv[1];
    int width, height;
    std::vector<unsigned char> rgbData = readPNGtoRGB(filename, width, height);
    std::vector<unsigned char> yuvData, rgbOut;
    if (!rgbData.empty()) {
        std::cout << "PNG converted to RGB successfully. Dimensions: " << width << "x" << height << std::endl;
        // 这里可以添加代码来处理或保存转换后的RGB数据
    } else {
        std::cerr << "Failed to convert PNG to RGB." << std::endl;
        return 1;
    }
    yuvData.resize(width * height * 3 / 2);

    TIMESTATIS_START(rgb2yuv);
    cvtRGBToNV21(rgbData.data(), yuvData.data(), width, height);
    TIMESTATIS_END(rgb2yuv);

    cv::Mat yuvY(height, width, CV_8UC1, yuvData.data());
    cv::Mat yuvUV(height / 2, width / 2, CV_8UC2, yuvData.data() + width * height);
    cv::Mat test(height, width, CV_8UC3);
    printf("%p\n", test.data);
    cv::Mat bgrImage;
    cv::cvtColorTwoPlane(yuvY, yuvUV, bgrImage, cv::COLOR_YUV2BGR_NV21);
    test.data = bgrImage.data;
    printf("%p %p\n", test.data, bgrImage.data);
    cv::imwrite("test1.jpg", bgrImage);
    cv::cvtColor(bgrImage, bgrImage, cv::COLOR_RGB2BGR);
    cv::imwrite("brgImage.jpg", bgrImage);

    test.data = bgrImage.data;

    printf("%p %p\n", test.data, bgrImage.data);
    cv::imwrite("test.jpg", bgrImage);

    uint8_t* p;
    int i = 10;
    while (i--) {
        TIMESTATIS_START(yuv2rgb);
        cvtNV21ToRGB(yuvData.data(), p, width, height);
        TIMESTATIS_END(yuv2rgb);
    }

    cv::Mat rgbo(height, width, CV_8UC3, p);
    cv::cvtColor(rgbo, rgbo, cv::COLOR_RGB2BGR);
    cv::imwrite("rgbOut.jpg", rgbo);

    cv::Mat bgr = cv::imread(argv[1]);
    cv::Mat yuv, bgrOut;

    TIMESTATIS_START(CVrgb2yuv);
    cv::cvtColor(bgr, yuv, cv::COLOR_BGR2YUV_I420);
    TIMESTATIS_END(CVrgb2yuv);

    TIMESTATIS_START(CVyuv2rgb);
    cv::cvtColor(yuv, bgrOut, cv::COLOR_YUV2BGR_I420);
    TIMESTATIS_END(CVyuv2rgb);

    cv::imwrite("bgrout.jpg", bgrOut);

    return 0;
}